﻿using System;
using System.Collections.Generic;
using Allegro.MathInterface;

namespace Allegro.Mathlib
{
    public class ParametrizedFunction : IParametricFunction
    {
        private ParametricFunction _func;

        public ParametrizedFunction(ParametricFunction func)
        {
            _func = func;
        }
        #region Implementation of IParametricFunction

        public double[] Parameters { get; set; }

        public double Value(double x)
        {
            return _func(x, Parameters);
        }

        #endregion
    }
    public class ChiSquareCurveFitting
    {
        private double[] _x;
        private double[] _y;
        private double[] _sigma;
        private IParametricFunction _parametricFunction;
        
        public ChiSquareCurveFitting(double[] x, double[] y, double[] sigma, IParametricFunction parametricFunction, double tolerance)
        {
            _x = x;
            _y = y;
            _sigma = sigma;
            _parametricFunction = parametricFunction;
            Tolerance = tolerance;
        }
        public ChiSquareCurveFitting(double[] x, double[] y, double[] sigma, ParametricFunction parametricFunction, double tolerance)
            : this(x, y, sigma, new ParametrizedFunction(parametricFunction), tolerance)
        {

        }

        public double Tolerance;

        public double NumIterations;

        /// <summary>
        /// Compute the chi-square of the model when fitted to the data
        /// </summary>
        /// <param name="parameters"></param>
        /// <returns></returns>
        public double ChiSquare(double[] parameters)
        {
            int l = _x.Length;
            double chi2 = 0.0;
            _parametricFunction.Parameters = parameters;
            Function f = _parametricFunction.Value;
            for (int i = 0; i < l; i++)
            {
                double yModel = f(_x[i]);
                double diff = (_y[i] - yModel) / _sigma[i];
                chi2 += diff * diff;
            }

            return chi2;
        }

        /// <summary>
        /// Find the parameters that minimizes the chi-square of the model when fitted to the data
        /// </summary>
        /// <param name="parameters"></param>
        /// <returns></returns>
        public double FindMinimum(double[] parameters)
        {
            VectorFunction f = ChiSquare;
            int dim = parameters.Length;
            double[] steps = new double[dim];
            for(int i = 0; i < dim; i++)
            {
                steps[i] = parameters[i]/10.0; // Heuristic choice of initial step
            }
            Gradient gradient = new Gradient(f, steps);
            ConjugateGradientMethod optimizer = new ConjugateGradientMethod(dim, f, gradient, Tolerance);

            double chi2Min = optimizer.FindMinimum(parameters);
            NumIterations = optimizer.NumIterations;
            return chi2Min;
        }
    }
}
